import arrow
import traceback

from pymongo import MongoClient, InsertOne, UpdateOne, DeleteOne

from conf.settings import CLOUDMONITOR_MONGODB

from utility.log import logger
from utility.utils import rollback, backup

conn = MongoClient(CLOUDMONITOR_MONGODB)

db = conn['cloudmonitor-dev5']

now_datetime = arrow.utcnow().datetime
now = now_datetime.isoformat()


def create_node_charges(o_uid, total, create_time, expire_time):
    _http = {
        'used': 0,
        'total': total,
        'limitType': 'frequency',
        'time': 0,
        'taskType': 'http',
        'funcType': 'httpMaxNode',
        'uid': o_uid,
        'chargeType': 'plan',
        'number': 1,
        'isDelete': False,
        'updateTime': now_datetime,
        'createTime': create_time,
        'expireTime': expire_time
    }
    _ping = {
        'used': 0,
        'total': total,
        'limitType': 'frequency',
        'time': 0,
        'taskType': 'ping',
        'funcType': 'pingMaxNode',
        'chargeType': 'plan',
        'number': 1,
        'isDelete': False,
        'uid': o_uid,
        'updateTime': now_datetime,
        'createTime': create_time,
        'expireTime': expire_time
    }

    update_list = [
        InsertOne(_http),
        InsertOne(_ping),
        DeleteOne({'uid': o_uid, 'funcType': 'httpMaxNode', 'total': 0}),
        DeleteOne({'uid': o_uid, 'funcType': 'pingMaxNode', 'total': 0}),
    ]

    db['charges'].bulk_write(update_list)


def main():
    backup_collections = ('balances', )
    backup(backup_collections, logger, db, now)

    try:
        plan_key_tuple = ('scanv_basic', 'scanv_custom_a')

        for plan_key in plan_key_tuple:
            plan = db['packages'].find_one({'key': plan_key, 'chargeType': 'plan'})

            update_list = []
            for balances in db['balances'].find({'planId': plan['_id']}):
                update_list.append(UpdateOne({'_id': balances['_id']}, {'$set': {
                    'balance.httpMaxNode.total': 5, 'balance.pingMaxNode.total': 5}}))

                create_node_charges(
                    balances['uid'], 5, balances['planChargeTime'], balances['planExpireTime'])

            if update_list:
                db['balances'].bulk_write(update_list)

    except Exception as e:
        logger.error(e)
        traceback.print_exc()
        rollback(backup_collections, logger, db, now)


if __name__ == '__main__':
    main()
